import os
import pandas as pd


class DataOperation:
    # TODO 判断是什么文件类型
    def __init__(self, path, sheet=0):
        """

        :param path: 文件路径
        :param sheet: excel表格
        """
        #  TODO 准备数据
        # 获取文件后缀
        _, extension = os.path.splitext(path)

        # 准备可支持的文件后缀
        extension_csv = ('.csv', '.tsv', 'txt')
        extension_excel = ('.xls', '.xlsx')

        # TODO 判断
        if extension in extension_csv:
            self.data = pd.read_csv(path)

        elif extension in extension_excel:
            self.data = pd.read_excel(path, sheet_name=sheet)

        else:
            print('不支持该类型的文件')

    # TODO 打印
    # 查看n条数据，默认是5条
    def head(self, item=5):
        print(self.data.head(item))

    # 从尾部查看,默认5条
    def tail(self, item=5):
        print(self.data.tail(item))

    # 查看数据形状
    def shape(self):
        print(self.data.shape)

    # 查看字段名
    def columns(self):
        print(self.data.columns)

    # 查看索引
    def index(self):
        print(self.data.index)

    # TODO 数据查询
    # 根据行列的标签值查询
    def show_loc(self, row_index=1, field=None):
        if field == None:
            print(self.data.loc[row_index].values)
        else:
            print(self.data.loc[row_index, field])

    # 根据行、列的数字位置查询
    def show_iloc(self, row_index=0, column_index=None):
        if column_index == None:
            print(self.data.iloc[row_index].values)
        else:
            print(self.data.iloc[row_index, column_index])

    # TODO 转化为字典
    def get_data_to_dict(self):
        """
        转化为字典
        根据标签查询到对应标签的所有值
        """
        return [self.data.loc[i].to_dict() for i in self.data.index.values]

    # TODO 转化为列表
    def get_data_to_list(self):
        """
        转化为列表
        """
        return self.data.values.tolist()


if __name__ == '__main__':
    pass
